package pir

import (
	"fmt"
	"math"
	"strings"
)

type State struct {
	Data []*Matrix
}
type Indexs struct {
	Data uint64
}
type Vec struct {
	Data []uint64
}

type CompressedState struct {
	Seed *PRGKey
}

type Msg struct {
	Data []*Matrix
}

// TMD，重构了。
func (m *Msg) Size() uint64 {
	sz := uint64(0)
	for _, d := range m.Data {
		sz += d.Size()
	}
	return sz
}

type MsgSlice struct {
	Data []Msg
}

func (m *MsgSlice) Size() uint64 {
	sz := uint64(0)
	for _, d := range m.Data {
		sz += d.Size()
	}
	return sz
}

func MakeState(elems ...*Matrix) State {
	st := State{}
	for _, elem := range elems {
		st.Data = append(st.Data, elem)
	}
	return st
}

func MakeCompressedState(elem *PRGKey) CompressedState {
	st := CompressedState{}
	st.Seed = elem
	return st
}

// 这个意思应该是拼接起来吧，message嘛
func MakeMsg(elems ...*Matrix) Msg {
	msg := Msg{}
	for _, elem := range elems {
		//fmt.Println("k")
		msg.Data = append(msg.Data, elem)
	}
	return msg
}

func MakeMsgSlice(elems ...Msg) MsgSlice {
	slice := MsgSlice{}
	for _, elem := range elems {
		slice.Data = append(slice.Data, elem)
	}
	return slice
}

// Returns the i-th elem in the representation of m in base p.
func Base_p(p, m, i uint64) uint64 {
	for j := uint64(0); j < i; j++ {
		m = m / p
	}
	return (m % p)
}

// Returns the element whose base-p decomposition is given by the values in vals
// p进制合并
func Reconstruct_from_base_p(p uint64, vals []uint64) uint64 {
	res := uint64(0)
	coeff := uint64(1)
	for _, v := range vals {
		res += coeff * v
		coeff *= p
	}
	return res
}

// Returns how many entries in Z_p are needed to represent an element in Z_q
func Compute_num_entries_base_p(p, log_q uint64) uint64 {
	log_p := uint64(math.Log2(float64(p)))
	return uint64(math.Ceil(float64(log_q) / float64(log_p)))
}

// Returns how many Z_p elements are needed to represent a database of N entries,
// each consisting of row_length bits.
func Num_DB_entries(N, row_length, p uint64) (uint64, uint64, uint64) {
	/*
		if float64(row_length) <= math.Log2(float64(p)) {
			// pack multiple DB entries into a single Z_p elem
			logp := uint64(math.Log2(float64(p)))
			entries_per_elem := logp / row_length
			db_entries := uint64(math.Ceil(float64(N) / float64(entries_per_elem)))
			if db_entries == 0 || db_entries > N {
				fmt.Printf("Num entries is %d; N is %d\n", db_entries, N)
				panic("Should not happen")
			}
			return db_entries, 1, entries_per_elem
		}
	*/
	// use multiple Z_p elems to represent a single DB entry
	ne := Compute_num_entries_base_p(p, row_length)
	return N * ne, ne, 0
}
func Num_DB_entries_with_long_text(N, row_length, p uint64) (uint64, uint64, uint64) {
	ne := Compute_num_entries_base_p(p, row_length)
	//fmt.Println("ne: ", ne)
	return N * ne, ne, 0
}
func avg(data []float64) float64 {
	sum := 0.0
	num := 0.0
	for _, elem := range data {
		sum += elem
		num += 1.0
	}
	return sum / num
}

func stddev(data []float64) float64 {
	avg := avg(data)
	sum := 0.0
	num := 0.0
	for _, elem := range data {
		sum += math.Pow(elem-avg, 2)
		num += 1.0
	}
	variance := sum / num // not -1!
	return math.Sqrt(variance)
}

func StringToBitString(s string) string {
	var builder strings.Builder
	for i := 0; i < len(s); i++ {
		// 逐字节处理，将每个字节转成8位二进制字符串
		builder.WriteString(fmt.Sprintf("%08b", s[i]))
	}
	return builder.String()
}

func StringToUint64ByBits(s string, bitLen uint64) []uint64 {
	if bitLen == 0 || bitLen > 64 {
		panic("bitLen must be between 1 and 64")
	}
	data := []byte(s)
	//totalBits := len(data) * 8
	var result []uint64

	var bitBuffer uint64 = 0
	var bitCount uint64 = 0

	for _, b := range data {
		bitBuffer = (bitBuffer << 8) | uint64(b)
		bitCount += 8

		for bitCount >= bitLen {
			shift := bitCount - bitLen
			val := (bitBuffer >> shift) & ((1 << bitLen) - 1)
			result = append(result, val)
			bitCount -= bitLen
			bitBuffer &= (1 << bitCount) - 1
		}
	}

	// 剩余不足bitLen的比特，作为最后一个块输出（低位填充）
	if bitCount > 0 {
		val := bitBuffer << (bitLen - bitCount)
		result = append(result, val&((1<<bitLen)-1))
	}

	return result
}

func Uint64ToStringByBits(data []uint64, bitLen uint64) string {
	if bitLen == 0 || bitLen > 64 {
		panic("bitLen must be between 1 and 64")
	}

	var bitBuffer uint64 = 0
	var bitCount uint64 = 0
	var bytes []byte

	for _, val := range data {
		// 将每个uint64值拼到bitBuffer后面
		bitBuffer = (bitBuffer << bitLen) | (val & ((1 << bitLen) - 1))
		bitCount += bitLen

		// 从bitBuffer高位取出8位，转成字节，直到不足8位
		for bitCount >= 8 {
			shift := bitCount - 8
			b := byte((bitBuffer >> shift) & 0xFF)
			bytes = append(bytes, b)
			bitCount -= 8
			bitBuffer &= (1 << bitCount) - 1 // 清除已经读过的高位字节
		}
	}
	// 剩余bitCount不足8位的部分丢弃（或根据需求保留）

	return string(bytes)
}
